package MapAndSet.Set;
import java.util.*;

public class TopKCountWords {
    public List<String> topKFrequent(String[] words, int k) {

        Map<String,Integer> map = new HashMap<>();
        for(String word:words) {
            if(map.get(word)==null) {
                map.put(word,1);
            }else {
                int val = map.get(word);
                map.put(word,val+1);
            }
        }

        PriorityQueue<Map.Entry<String,Integer>> minHeap = new PriorityQueue<>(new Comparator<Map.Entry<String, Integer>>() {
            @Override
            public int compare(Map.Entry<String, Integer> o1, Map.Entry<String, Integer> o2) {

                if(o1.getValue().compareTo(o2.getValue())==0) {
                    return o2.getKey().compareTo(o1.getKey());
                }

                return o1.getValue()- o2.getValue();
            }
        });

        for(Map.Entry<String,Integer> entry:map.entrySet()) {
            if(minHeap.size() < k) {
                minHeap.offer(entry);
            }else {
                Map.Entry<String,Integer> top = minHeap.peek();
                if(top.getValue().equals(entry.getValue())) {
                    if(top.getKey().compareTo(entry.getKey())>0) {
                        minHeap.poll();
                        minHeap.offer(entry);
                    }
                } else {
                    if(top.getValue().compareTo(entry.getValue())<0) {
                        minHeap.poll();
                        minHeap.offer(entry);
                    }
                }
            }
        }

        List<String> ret = new ArrayList<>();
        for(int i = 0 ; i < k; i++ ) {
            Map.Entry<String,Integer> top = minHeap.poll();
            ret.add(top.getKey());
        }
        Collections.reverse(ret);

        return ret;
    }
}
